import torch
from lib.similarity import euclid_dist
from lib.base import BaseMetric


# calculate the acc by the Euclidean distance
# @METRIC.register('protonet_acc')
class ProtoNetAcc(BaseMetric):
    def __init__(self, cfg):
        super().__init__(cfg)
        self.cfg = cfg
        self.similarity = euclid_dist

    def __call__(self, data_f, labels, type):
        # similarity = self.similarity[type]

        batch_size = self.batch_size[type]
        n_way = self.n_way[type]
        n_query = self.n_query[type]
        total_query = n_way * n_query  # 300

        support_sets_f = data_f[:, :n_way, :]
        query_sets_f = data_f[:, n_way:, :]

        # sim_matrix: 10*300*60
        sim_matrix = self.similarity(support_sets_f, query_sets_f)
        # sim_matrix: 3000*60
        # sim_matrix = sim_matrix.flatten(start_dim=0, end_dim=1)

        # notice: the result of acc should be the mean of batch_size number episodes,
        # and shouldn't be the correct number divides total number

        labels = torch.arange(n_way).to(labels.device)
        # labels: 60 -> 1*60*1 -> 10*60*5
        labels = labels.view(1, n_way, 1).expand(batch_size, n_way, n_query)
        # labels: 10*60*5 -> 10*300
        labels = labels.flatten(start_dim=1, end_dim=2)
        # sim_matrix is Euclidean distance!!!!!!!!!! The smaller, the more similar!!!!!!!!!!!!!!
        # preds: 10*300
        preds = torch.argmax(sim_matrix, dim=-1)
        # acc: 10*1
        acc = (labels == preds).float().sum(dim=1) / total_query
        # acc: 1*1
        acc = acc.mean()

        return acc
